#include <ros/ros.h>
#include <sensor_msgs/Image.h>
#include <cv_bridge/cv_bridge.h>
#include <image_transport/image_transport.h>

#include <NvInfer.h>
#include <cuda_runtime_api.h>

#include <opencv2/opencv.hpp>

#include <algorithm>
#include <chrono>
#include <memory>
#include <string>
#include <vector>

class TrtLogger : public nvinfer1::ILogger {
public:
    void log(nvinfer1::ILogger::Severity severity, const char* msg) noexcept override {
        if (severity <= nvinfer1::ILogger::Severity::kWARNING) {
            ROS_WARN("[TRT] %s", msg);
        }
    }
};

class DeepLabTrtNode {
private:
    ros::NodeHandle nh_;
    image_transport::ImageTransport it_;
    image_transport::Subscriber image_sub_;
    image_transport::Publisher overlay_pub_;
    
    // TensorRT components
    std::unique_ptr<nvinfer1::IRuntime, void(*)(nvinfer1::IRuntime*)> runtime_;
    std::unique_ptr<nvinfer1::ICudaEngine, void(*)(nvinfer1::ICudaEngine*)> engine_;
    std::unique_ptr<nvinfer1::IExecutionContext, void(*)(nvinfer1::IExecutionContext*)> context_;
    
    // CUDA buffers
    void* d_input_;
    void* d_output_;
    
    // Model parameters
    static constexpr int TARGET_H = 1024;
    static constexpr int TARGET_W = 2048;
    static constexpr int INPUT_C = 3;
    static constexpr int NUM_CLASSES = 19;
    
    // Binding indices
    int input_index_;
    int output_index_;
    
    bool initialized_;

public:
    DeepLabTrtNode() : nh_("~"), it_(nh_), initialized_(false) {
        // Initialize TensorRT
        if (!initializeTensorRT()) {
            ROS_ERROR("Failed to initialize TensorRT");
            return;
        }
        
        // Setup ROS
        image_sub_ = it_.subscribe("input_image", 1, &DeepLabTrtNode::imageCallback, this);
        overlay_pub_ = it_.advertise("overlay_image", 1);
        
        ROS_INFO("DeepLab TensorRT ROS node initialized");
        ROS_INFO("Subscribing to: %s", image_sub_.getTopic().c_str());
        ROS_INFO("Publishing to: %s", overlay_pub_.getTopic().c_str());
    }
    
    ~DeepLabTrtNode() {
        if (d_input_) cudaFree(d_input_);
        if (d_output_) cudaFree(d_output_);
    }
    
    bool initializeTensorRT() {
        std::string engine_path;
        nh_.param<std::string>("engine_path", engine_path, "checkpoints/seg.trt");
        
        // Load engine file
        std::ifstream file(engine_path, std::ios::binary | std::ios::ate);
        if (!file) {
            ROS_ERROR("Cannot open engine file: %s", engine_path.c_str());
            return false;
        }
        
        std::streamsize size = file.tellg();
        file.seekg(0, std::ios::beg);
        std::vector<char> engine_data(size);
        if (!file.read(engine_data.data(), size)) {
            ROS_ERROR("Failed to read engine file");
            return false;
        }
        
        // Create runtime and engine
        TrtLogger logger;
        runtime_.reset(nvinfer1::createInferRuntime(logger), [](nvinfer1::IRuntime* p){ if(p) p->destroy(); });
        if (!runtime_) {
            ROS_ERROR("Failed to create TensorRT runtime");
            return false;
        }
        
        engine_.reset(runtime_->deserializeCudaEngine(engine_data.data(), engine_data.size()),
                     [](nvinfer1::ICudaEngine* p){ if(p) p->destroy(); });
        if (!engine_) {
            ROS_ERROR("Failed to deserialize engine");
            return false;
        }
        
        context_.reset(engine_->createExecutionContext(), [](nvinfer1::IExecutionContext* p){ if(p) p->destroy(); });
        if (!context_) {
            ROS_ERROR("Failed to create execution context");
            return false;
        }
        
        // Find input/output bindings
        int nb_bindings = engine_->getNbBindings();
        for (int i = 0; i < nb_bindings; ++i) {
            if (engine_->bindingIsInput(i)) {
                input_index_ = i;
            } else {
                output_index_ = i;
            }
        }
        
        // Set binding dimensions
        nvinfer1::Dims input_dims = engine_->getBindingDimensions(input_index_);
        if (input_dims.nbDims == 3 || input_dims.nbDims == 4) {
            nvinfer1::Dims4 d(1, INPUT_C, TARGET_H, TARGET_W);
            context_->setBindingDimensions(input_index_, d);
        }
        
        // Allocate CUDA buffers
        size_t input_size = 1 * INPUT_C * TARGET_H * TARGET_W * sizeof(float);
        size_t output_size = 1 * NUM_CLASSES * TARGET_H * TARGET_W * sizeof(float);
        
        cudaError_t err;
        err = cudaMalloc(&d_input_, input_size);
        if (err != cudaSuccess) {
            ROS_ERROR("Failed to allocate input buffer: %s", cudaGetErrorString(err));
            return false;
        }
        
        err = cudaMalloc(&d_output_, output_size);
        if (err != cudaSuccess) {
            ROS_ERROR("Failed to allocate output buffer: %s", cudaGetErrorString(err));
            return false;
        }
        
        initialized_ = true;
        ROS_INFO("TensorRT engine loaded successfully");
        return true;
    }
    
    void imageCallback(const sensor_msgs::ImageConstPtr& msg) {
        if (!initialized_) {
            ROS_WARN("TensorRT not initialized, skipping frame");
            return;
        }
        
        try {
            // Convert ROS image to OpenCV
            cv_bridge::CvImagePtr cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
            cv::Mat original_bgr = cv_ptr->image;
            
            // Preprocess
            auto t0 = std::chrono::steady_clock::now();
            cv::Mat resized_bgr;
            cv::resize(original_bgr, resized_bgr, cv::Size(TARGET_W, TARGET_H), 0, 0, cv::INTER_LINEAR);
            
            cv::Mat rgb;
            cv::cvtColor(resized_bgr, rgb, cv::COLOR_BGR2RGB);
            cv::Mat rgb_f;
            rgb.convertTo(rgb_f, CV_32FC3, 1.0 / 255.0);
            
            // Normalize with ImageNet stats
            const float mean[3] = {0.485f, 0.456f, 0.406f};
            const float std[3] = {0.229f, 0.224f, 0.225f};
            
            std::vector<float> input_host(INPUT_C * TARGET_H * TARGET_W);
            for (int y = 0; y < TARGET_H; ++y) {
                const cv::Vec3f* pr = rgb_f.ptr<cv::Vec3f>(y);
                for (int x = 0; x < TARGET_W; ++x) {
                    for (int c = 0; c < INPUT_C; ++c) {
                        float v = (pr[x][c] - mean[c]) / std[c];
                        input_host[c * TARGET_H * TARGET_W + y * TARGET_W + x] = v;
                    }
                }
            }
            auto t1 = std::chrono::steady_clock::now();
            
            // Copy to GPU and run inference
            cudaMemcpy(d_input_, input_host.data(), input_host.size() * sizeof(float), cudaMemcpyHostToDevice);
            
            void* bindings[2];
            bindings[input_index_] = d_input_;
            bindings[output_index_] = d_output_;
            
            auto t2 = std::chrono::steady_clock::now();
            if (!context_->enqueueV2(bindings, nullptr, nullptr)) {
                ROS_ERROR("Inference failed");
                return;
            }
            cudaDeviceSynchronize();
            auto t3 = std::chrono::steady_clock::now();
            
            // Copy output back
            std::vector<float> output_host(NUM_CLASSES * TARGET_H * TARGET_W);
            cudaMemcpy(output_host.data(), d_output_, output_host.size() * sizeof(float), cudaMemcpyDeviceToHost);
            
            // Postprocess
            auto t4 = std::chrono::steady_clock::now();
            cv::Mat mask(TARGET_H, TARGET_W, CV_8UC1);
            for (int y = 0; y < TARGET_H; ++y) {
                uint8_t* pm = mask.ptr<uint8_t>(y);
                for (int x = 0; x < TARGET_W; ++x) {
                    int idx_base = y * TARGET_W + x;
                    int best_id = 0;
                    float best_val = output_host[idx_base];
                    for (int c = 1; c < NUM_CLASSES; ++c) {
                        float v = output_host[c * TARGET_H * TARGET_W + idx_base];
                        if (v > best_val) { best_val = v; best_id = c; }
                    }
                    pm[x] = static_cast<uint8_t>(best_id);
                }
            }
            
            // Colorize mask
            cv::Mat color_rgb = colorizeMask(mask);
            
            // Resize to original size
            cv::Mat color_rgb_orig;
            cv::resize(color_rgb, color_rgb_orig, cv::Size(original_bgr.cols, original_bgr.rows), 0, 0, cv::INTER_NEAREST);
            
            // Create overlay
            cv::Mat color_bgr_orig;
            cv::cvtColor(color_rgb_orig, color_bgr_orig, cv::COLOR_RGB2BGR);
            cv::Mat overlay;
            const double alpha = 0.7;
            cv::addWeighted(original_bgr, 1.0 - alpha, color_bgr_orig, alpha, 0.0, overlay);
            auto t5 = std::chrono::steady_clock::now();
            
            // Publish overlay result
            cv_bridge::CvImage overlay_msg;
            overlay_msg.header = msg->header;
            overlay_msg.encoding = sensor_msgs::image_encodings::BGR8;
            overlay_msg.image = overlay;
            overlay_pub_.publish(overlay_msg.toImageMsg());
            
            // Log timing
            auto dur_prep = std::chrono::duration_cast<std::chrono::milliseconds>(t1 - t0).count();
            auto dur_inf = std::chrono::duration_cast<std::chrono::milliseconds>(t3 - t2).count();
            auto dur_post = std::chrono::duration_cast<std::chrono::milliseconds>(t5 - t4).count();
            
            ROS_DEBUG("Timing - Prep: %.1fms, Inf: %.1fms, Post: %.1fms", 
                     dur_prep/1000.0, dur_inf/1000.0, dur_post/1000.0);
            
        } catch (cv_bridge::Exception& e) {
            ROS_ERROR("cv_bridge exception: %s", e.what());
        }
    }
    
    cv::Mat colorizeMask(const cv::Mat& mask) {
        // Cityscapes colormap (19 classes)
        static const cv::Vec3b colors[19] = {
            {128, 64, 128}, {244, 35, 232}, {70, 70, 70}, {102, 102, 156}, {190, 153, 153},
            {153, 153, 153}, {250, 170, 30}, {220, 220, 0}, {107, 142, 35}, {152, 251, 152},
            {70, 130, 180}, {220, 20, 60}, {255, 0, 0}, {0, 0, 142}, {0, 0, 70},
            {0, 60, 100}, {0, 80, 100}, {0, 0, 230}, {119, 11, 32}
        };
        
        cv::Mat color(mask.size(), CV_8UC3);
        for (int y = 0; y < mask.rows; ++y) {
            const uint8_t* pm = mask.ptr<uint8_t>(y);
            cv::Vec3b* pc = color.ptr<cv::Vec3b>(y);
            for (int x = 0; x < mask.cols; ++x) {
                uint8_t id = pm[x];
                if (id < 19) pc[x] = colors[id];
                else pc[x] = {0, 0, 0};
            }
        }
        return color;
    }
};

int main(int argc, char** argv) {
    ros::init(argc, argv, "deeplab_trt_node");
    
    try {
        DeepLabTrtNode node;
        ros::spin();
    } catch (const std::exception& e) {
        ROS_ERROR("Exception: %s", e.what());
        return 1;
    }
    
    return 0;
}
